import pandas as pd
import numpy as np
import nibabel as nib
import cv2
from sklearn.model_selection import train_test_split
import SimpleITK as sitk
import random


def pprint(sitkImage):
    print("原点位置:{}".format(sitkImage.GetOrigin()))
    print("尺寸：{}".format(sitkImage.GetSize()))
    print("体素大小(x,y,z):{}".format(sitkImage.GetSpacing()))
    print("图像方向:{}".format(sitkImage.GetDirection()))

    # 查看图像相关的纬度信息
    print("维度:{}".format(sitkImage.GetDimension()))
    print("宽度:{}".format(sitkImage.GetWidth()))
    print("高度:{}".format(sitkImage.GetHeight()))
    print("深度(层数):{}".format(sitkImage.GetDepth()))


def sitk_read_raw(img_path, mode='train', L=-600, W=1500):
    # 读取nifit格式，设置窗口窗位W=1500，L=-600，并缩放到0到255之间，并转换成numpy格式,针对肺部
    sitkImage = sitk.ReadImage(img_path)
    intensityWindowingFilter = sitk.IntensityWindowingImageFilter()
    # 转换成0到255之间
    # intensityWindowingFilter.SetOutputMinimum(0)
    # intensityWindowingFilter.SetOutputMaximum(255)
    min = int(L - W / 2.0)
    max = int(L + W / 2.0)
    print('min', min, 'max', max)
    if 'mask' not in img_path:
        # 调窗宽窗位
        # 相当于数据增强中的 改变对比度
        if mode == 'train':
            max_num = random.randint(1000, 1900)
            min_num = random.randint(-900, -300)
            intensityWindowingFilter.SetWindowMaximum(max_num)
            intensityWindowingFilter.SetWindowMinimum(min_num)
        else:
            intensityWindowingFilter.SetWindowMaximum(max)
            intensityWindowingFilter.SetWindowMinimum(min)
    # 设置窗宽窗位，并进行缩放
    sitkImage = intensityWindowingFilter.Execute(sitkImage)
    nda = sitk.GetArrayFromImage(sitkImage)
    return np.transpose(nda, (1, 2, 0))


def Scale(sitkImage):
    # 通过RescaleIntensity方法既可以进行缩放，没有窗位的调整
    return sitk.RescaleIntensity(sitkImage)


def read_nii(filepath):
    '''
    Reads .nii file and returns pixel array
    '''
    ct_scan = nib.load(filepath)
    array = ct_scan.get_fdata()
    array = np.rot90(np.array(array))
    print(array.shape)
    return array


def generate(ct, mask, Img_size=128):
    # Read sample
    CT = []
    GT = []
    num = ct.shape[-1]

    for slice_id in range(1, num - 1):
        """取相邻的切片，2.5D"""
        vol = cv2.resize(ct[..., slice_id - 1:slice_id + 2], dsize=(Img_size, Img_size),
                         interpolation=cv2.INTER_AREA).astype('float32')
        lab = cv2.resize(mask[..., slice_id:slice_id + 1], dsize=(Img_size, Img_size),
                         interpolation=cv2.INTER_AREA).astype('float32')
        if np.sum(lab) <= 10:  # 去掉前景过少的数据
            continue
        CT.append(vol)
        GT.append(lab)
    CT = np.array(CT)
    GT = np.array(GT)
    return CT, GT


def clahe_equalized(imgs, start, end):
    assert (len(imgs.shape) == 3)  # 3D arrays
    # create a CLAHE object (Arguments are optional).
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    imgs_equalized = np.empty(imgs.shape)
    for i in range(start, end + 1):
        imgs_equalized[i, :, :] = clahe.apply(np.array(imgs[i, :, :], dtype=np.uint8))
    return imgs_equalized


def batch_func(x_train, y_train, Img_size=128):
    CT = None
    GT = None
    for x, y in zip(x_train, y_train):
        ct = sitk_read_raw(x, mode='valid')
        gt = sitk_read_raw(y, mode='valid')
        ct, gt = generate(ct, gt, Img_size=Img_size)
        if CT is not None:
            print('CT', CT.shape)
            print('ct', ct.shape)
            print('GT', GT.shape)
            print('gt', gt.shape)
            CT = np.concatenate([CT, ct], axis=0)
            GT = np.concatenate([GT, gt], axis=0)
        else:
            CT = ct
            GT = gt

    return CT, GT


def clip_label(label, category):
    '''针对有多类标签的情况，将category和比category大的类别赋为1,其余的赋为0
    label: 分割的标签
    category: 保留的类别
    '''
    label[label < category] = 0
    label[label >= category] = 1
    return label


def ImagScale(img_path):
    # 通过RescaleIntensity方法既可以进行缩放，没有窗位的调整
    sitkImage = sitk.ReadImage(img_path)
    sitkImage = sitk.RescaleIntensity(sitkImage)
    nda = sitk.GetArrayFromImage(sitkImage)
    return np.transpose(nda, (1, 2, 0))


if __name__ == '__main__':
    root = '/hdd9/ppp/ImageSegments/COVID'

    raw_data = pd.read_csv(root + '/metadata.csv')
    for cc in raw_data.columns:
        raw_data[cc] = raw_data[cc].apply(lambda x: root + '/' + '/'.join(x.split('/')[-2:]))

    for name in ['lung_mask']:
        print('processing {}'.format(name))
        category = 1
        sample_ct = None
        sample_gt = None
        ct_scan = raw_data['ct_scan'].tolist()
        mask = raw_data[name].tolist()
        # x_train, x_test, y_train, x_test = train_test_split(ct_scan, mask, test_size=0.1)
        sample_ct, sample_gt = batch_func(ct_scan, mask, 128)
        print('min={},max={}'.format(np.min(sample_gt), np.max(sample_gt)))
        sample_gt = sample_gt[..., np.newaxis]
        sample_gt = clip_label(sample_gt, category=category)
        print('min={},max={}'.format(np.min(sample_gt), np.max(sample_gt)))
        print(sample_ct.shape)
        print(sample_gt.shape)
        x_train, x_test, y_train, y_test = train_test_split(sample_ct, sample_gt, test_size=0.1)
        np.save(root + '/train_ct_{}.npy'.format(name), x_train)
        np.save(root + '/train_gt_{}.npy'.format(name), y_train)

        '''test'''
        np.save(root + '/test_ct_{}.npy'.format(name), x_test)
        np.save(root + '/test_gt_{}.npy'.format(name), y_test)
